package com.superai.controller;

import io.github.imfangs.dify.client.DifyWorkflowClient;
import io.github.imfangs.dify.client.callback.WorkflowStreamCallback;
import io.github.imfangs.dify.client.enums.ResponseMode;
import io.github.imfangs.dify.client.event.*;
import io.github.imfangs.dify.client.model.workflow.*;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 工作流应用控制器
 *
 * 这个控制器提供了与Dify工作流应用交互的API，包括：
 * - 执行工作流（阻塞模式和流式模式）
 * - 停止工作流
 * - 获取工作流执行状态
 * - 获取工作流日志
 */
@RestController
@RequestMapping("/workflow")
@RequiredArgsConstructor
public class WorkflowController {

    /**
     * Dify工作流客户端，通过构造函数注入
     */
    @Autowired
    private  DifyWorkflowClient workflowClient;

    /**
     * 存储SSE连接的映射表，用于流式响应
     */
    private final ConcurrentHashMap<String, SseEmitter> sseEmitters = new ConcurrentHashMap<>();

    /**
     * 执行工作流（阻塞模式）
     *
     * @param request 工作流执行请求
     * @return 工作流执行响应
     */
    @PostMapping("/run")
    public ResponseEntity<WorkflowRunResponse> runWorkflow(@RequestBody WorkflowRequestDTO request) {
        try {
            // 创建工作流请求
            WorkflowRunRequest workflowRequest = WorkflowRunRequest.builder()
                    .inputs(request.getInputs())
                    .user(request.getUser())
                    .responseMode(ResponseMode.BLOCKING)
                    .build();

            // 执行工作流并获取响应
            WorkflowRunResponse response = workflowClient.runWorkflow(workflowRequest);
            return ResponseEntity.ok(response);
        } catch (IOException e) {
            // 记录异常并返回错误响应
            return ResponseEntity.internalServerError().build();
        }
    }

    /**
     * 执行工作流（流式模式）
     *
     * @param request 工作流执行请求
     * @return SSE发射器，用于流式传输响应
     */
    @PostMapping(value = "/run/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter runWorkflowStream(@RequestBody WorkflowRequestDTO request) {
        // 创建SSE发射器，设置超时时间为30分钟
        SseEmitter emitter = new SseEmitter(1800000L);
        String emitterId = request.getUser() + "_" + System.currentTimeMillis();
        sseEmitters.put(emitterId, emitter);

        try {
            // 创建工作流请求
            WorkflowRunRequest workflowRequest = WorkflowRunRequest.builder()
                    .inputs(request.getInputs())
                    .user(request.getUser())
                    .responseMode(ResponseMode.STREAMING)
                    .build();

            // 执行工作流流式请求
            workflowClient.runWorkflowStream(workflowRequest, new WorkflowStreamCallback() {
                @Override
                public void onWorkflowStarted(WorkflowStartedEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("workflow_started")
                                .data(event));
                    } catch (IOException e) {
                        onException(e);
                    }
                }

                @Override
                public void onNodeStarted(NodeStartedEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("node_started")
                                .data(event));
                    } catch (IOException e) {
                        onException(e);
                    }
                }

                @Override
                public void onNodeFinished(NodeFinishedEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("node_finished")
                                .data(event));
                    } catch (IOException e) {
                        onException(e);
                    }
                }

                @Override
                public void onWorkflowFinished(WorkflowFinishedEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("workflow_finished")
                                .data(event));
                        emitter.complete();
                    } catch (IOException e) {
                        onException(e);
                    } finally {
                        sseEmitters.remove(emitterId);
                    }
                }

                @Override
                public void onError(ErrorEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("error")
                                .data(event));
                        emitter.complete();
                    } catch (IOException e) {
                        onException(e);
                    } finally {
                        sseEmitters.remove(emitterId);
                    }
                }

                @Override
                public void onException(Throwable throwable) {
                    emitter.completeWithError(throwable);
                    sseEmitters.remove(emitterId);
                }
            });
        } catch (IOException e) {
            emitter.completeWithError(e);
            sseEmitters.remove(emitterId);
        }

        // 设置超时和完成回调
        emitter.onTimeout(() -> sseEmitters.remove(emitterId));
        emitter.onCompletion(() -> sseEmitters.remove(emitterId));

        return emitter;
    }

    /**
     * 停止工作流
     *
     * @param taskId 任务ID
     * @param user 用户ID
     * @return 工作流停止响应
     */
    @PostMapping("/stop")
    public ResponseEntity<WorkflowStopResponse> stopWorkflow(
            @RequestParam String taskId,
            @RequestParam String user) {
        try {
            WorkflowStopResponse response = workflowClient.stopWorkflow(taskId, user);
            return ResponseEntity.ok(response);
        } catch (IOException e) {
            // 记录异常并返回错误响应
            return ResponseEntity.internalServerError().build();
        }
    }

    /**
     * 获取工作流执行状态
     *
     * @param workflowRunId 工作流执行ID
     * @return 工作流执行状态响应
     */
    @GetMapping("/runs/{workflowRunId}")
    public ResponseEntity<WorkflowRunStatusResponse> getWorkflowRun(
            @PathVariable String workflowRunId) {
        try {
            WorkflowRunStatusResponse response = workflowClient.getWorkflowRun(workflowRunId);
            return ResponseEntity.ok(response);
        } catch (IOException e) {
            // 记录异常并返回错误响应
            return ResponseEntity.internalServerError().build();
        }
    }

    /**
     * 获取工作流日志
     *
     * @param keyword 关键词（可选）
     * @param status 状态（可选）
     * @param page 页码
     * @param limit 每页记录数
     * @return 工作流日志响应
     */
    @GetMapping("/logs")
    public ResponseEntity<WorkflowLogsResponse> getWorkflowLogs(
            @RequestParam(required = false) String keyword,
            @RequestParam(required = false) String status,
            @RequestParam(defaultValue = "1") Integer page,
            @RequestParam(defaultValue = "20") Integer limit) {
        try {
            WorkflowLogsResponse response = workflowClient.getWorkflowLogs(keyword, status, page, limit);
            return ResponseEntity.ok(response);
        } catch (IOException e) {
            // 记录异常并返回错误响应
            return ResponseEntity.internalServerError().build();
        }
    }

    /**
     * 工作流请求数据传输对象
     */
    public static class WorkflowRequestDTO {
        private Map<String, Object> inputs;
        private String user;

        // Getters and setters
        public Map<String, Object> getInputs() {
            return inputs;
        }

        public void setInputs(Map<String, Object> inputs) {
            this.inputs = inputs;
        }

        public String getUser() {
            return user;
        }

        public void setUser(String user) {
            this.user = user;
        }
    }
}
